{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Optimize Water Distribution in a Village"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #union-find #graph #minimum-spanning-tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #并查集 #图 #最小生成树"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: minCostToSupplyWater"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #水资源分配优化"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>村里面一共有 <code>n</code> 栋房子。我们希望通过建造水井和铺设管道来为所有房子供水。</p>\n",
    "\n",
    "<p>对于每个房子&nbsp;<code>i</code>，我们有两种可选的供水方案：一种是直接在房子内建造水井，成本为&nbsp;<code>wells[i - 1]</code>&nbsp;（注意 <code>-1</code> ，因为 <strong>索引从0开始</strong> ）；另一种是从另一口井铺设管道引水，数组&nbsp;<code>pipes</code>&nbsp;给出了在房子间铺设管道的成本，其中每个&nbsp;<code>pipes[j] = [house1<sub>j</sub>, house2<sub>j</sub>, cost<sub>j</sub>]</code>&nbsp;代表用管道将&nbsp;<code>house1<sub>j</sub></code>&nbsp;和&nbsp;<code>house2<sub>j</sub></code>连接在一起的成本。连接是双向的。</p>\n",
    "\n",
    "<p>请返回 <em>为所有房子都供水的最低总成本</em> 。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<p><strong><img alt=\"\" src=\"https://assets.leetcode-cn.com/aliyun-lc-upload/uploads/2019/08/23/1359_ex1.png\" /></strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]\n",
    "<strong>输出：</strong>3\n",
    "<strong>解释： </strong>\n",
    "上图展示了铺设管道连接房屋的成本。\n",
    "最好的策略是在第一个房子里建造水井（成本为 1），然后将其他房子铺设管道连起来（成本为 2），所以总成本为 3。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>n = 2, wells = [1,1], pipes = [[1,2,1]]\n",
    "<strong>输出：</strong>2\n",
    "<strong>解释：</strong>我们可以用以下三种方法中的一种来提供低成本的水:\n",
    "选项1:\n",
    "在1号房子里面建一口井，成本为1\n",
    "在房子2内建造井，成本为1\n",
    "总成本是2。\n",
    "选项2:\n",
    "在1号房子里面建一口井，成本为1\n",
    "-花费1连接房子2和房子1。\n",
    "总成本是2。\n",
    "选项3:\n",
    "在房子2内建造井，成本为1\n",
    "-花费1连接房子1和房子2。\n",
    "总成本是2。\n",
    "注意，我们可以用cost 1或cost 2连接房子1和房子2，但我们总是选择最便宜的选项。</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>2 &lt;= n &lt;= 10<sup>4</sup></code></li>\n",
    "\t<li><code>wells.length == n</code></li>\n",
    "\t<li><code>0 &lt;= wells[i] &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>1 &lt;= pipes.length &lt;= 10<sup>4</sup></code></li>\n",
    "\t<li><code>pipes[j].length == 3</code></li>\n",
    "\t<li><code>1 &lt;= house1<sub>j</sub>, house2<sub>j</sub>&nbsp;&lt;= n</code></li>\n",
    "\t<li><code>0 &lt;= cost<sub>j</sub>&nbsp;&lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>house1<sub>j</sub>&nbsp;!= house2<sub>j</sub></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [optimize-water-distribution-in-a-village](https://leetcode.cn/problems/optimize-water-distribution-in-a-village/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [optimize-water-distribution-in-a-village](https://leetcode.cn/problems/optimize-water-distribution-in-a-village/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['3\\n[1,2,2]\\n[[1,2,1],[2,3,1]]', '2\\n[1,1]\\n[[1,2,1],[1,2,2]]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import heapq\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        adjList = defaultdict(list)\n",
    "        for h1, h2, cost in pipes:\n",
    "            adjList[h1].append((cost, h2))\n",
    "            adjList[h2].append((cost, h1))\n",
    "        for i in range(n):\n",
    "            adjList[0].append((wells[i], i+1))\n",
    "        heap = [(0, 0)]\n",
    "        visited = set()\n",
    "        total = 0\n",
    "        while heap:\n",
    "            cost, house = heapq.heappop(heap)\n",
    "            if house not in visited:\n",
    "                visited.add(house)\n",
    "                total += cost\n",
    "                for cost, n_house in adjList[house]:\n",
    "                    if n_house not in visited:\n",
    "                        heapq.heappush(heap, (cost, n_house))\n",
    "            if len(visited) == n + 1:\n",
    "                break\n",
    "        return total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    '''\n",
    "    1. at least one house has well built inside house\n",
    "    2. all housed not need to be connected to get water supplied\n",
    "    3. the cost of building a well in current group is smaller than connecting current group to a already supplied group\n",
    "    '''\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        id = [i for i in range(n)] # the root of house is itself at the beginning\n",
    "        cost = [wells[i] for i in range(n)] # the cost of building well in current group\n",
    "        edges_cost = 0\n",
    "\n",
    "        pipes = sorted(pipes, key=lambda x: x[-1])  # Sort by the last element of each sublist\n",
    "\n",
    "\n",
    "        for h1, h2, c in pipes:\n",
    "            p = self.find_root(id, h1 - 1)\n",
    "            q = self.find_root(id, h2 - 1)\n",
    "            if p == q:\n",
    "                continue \n",
    "            cur_combine_cost = c + min(cost[p], cost[q])\n",
    "            cur_individual_cost = cost[p] + cost[q]\n",
    "            if cur_combine_cost <= cur_individual_cost:\n",
    "                id[p] = q\n",
    "                cost[q] = min(cost[q], cost[p])\n",
    "                edges_cost += c\n",
    "\n",
    "            # if  c > cost[h2-1] and c >  cost[h1-1]:\n",
    "            #     continue # the cost of building h1/h2 and connect to h1/h2 costs more than them build themselves\n",
    "            # # connect the lowest cost path\n",
    "            # if c <= cost[h2-1] and cost[h1-1] < cost[h2-1]: # h1 build wells\n",
    "            #     p = self.find_root(id, h1 - 1)\n",
    "            #     q = self.find_root(id, h2 - 1)\n",
    "            #     id[q] = p \n",
    "            #     cost[h2-1] = c \n",
    "            \n",
    "            # elif c <= cost[h1-1] and cost[h2-1] < cost[h1-1]: # h2 build wells\n",
    "            #     p = self.find_root(id, h2 - 1)\n",
    "            #     q = self.find_root(id, h1 - 1)\n",
    "            #     id[q] = p \n",
    "            #     cost[h1-1] = c        \n",
    "\n",
    "        res = 0 \n",
    "        for i in range(n):\n",
    "            if i == id[i]:\n",
    "                res += cost[i]\n",
    "\n",
    "        return res+edges_cost\n",
    "\n",
    "    def find_root(self, id, i):\n",
    "        while i != id[i]:\n",
    "            id[i] = id[id[i]]\n",
    "            i = id[i]\n",
    "        return i \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class DSU:\n",
    "    def __init__(self, n):\n",
    "        self.data = [i for i in range(n)]\n",
    "    def find(self, x):\n",
    "        if x == self.data[x]:\n",
    "            return x \n",
    "        else:\n",
    "            self.data[x] = self.find(self.data[x])\n",
    "            return self.data[x] \n",
    "    def un(self, x, y):\n",
    "        fx,fy = self.find(x), self.find(y)\n",
    "        if fx == fy:\n",
    "            return False \n",
    "        else:\n",
    "            self.data[fy] = fx \n",
    "            return True \n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        edges = [(wells[i], 0, i+1) for i in range(n)]\n",
    "        for a, b, w in pipes:\n",
    "            edges.append((w, a, b))\n",
    "        edges.sort() \n",
    "        dsu = DSU(n+1) \n",
    "        ans = 0 \n",
    "        #print(edges)\n",
    "        for w, a, b in edges:\n",
    "            if dsu.un(a, b):\n",
    "                ans += w\n",
    "        #print(dsu.data) \n",
    "        return ans \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        num_node = len(wells)\n",
    "        edges = []\n",
    "        for i in range(len(wells)):\n",
    "            edges.append([0,i+1,wells[i]])\n",
    "        \n",
    "        \n",
    "        edges.extend(pipes)\n",
    "        \n",
    "        edges.sort(key=lambda x:x[2])\n",
    "        fathers = [x for x in range(num_node+1)]\n",
    "\n",
    "        def find(ele):\n",
    "            if ele == fathers[ele]:\n",
    "                return ele\n",
    "\n",
    "            fathers[ele] = find(fathers[ele])\n",
    "            return fathers[ele]\n",
    "\n",
    "        used_edges = []\n",
    "        ans = 0\n",
    "        for i,j,cost in edges:\n",
    "            father_i = find(i)\n",
    "            father_j = find(j)\n",
    "            if father_i != father_j:\n",
    "                used_edges.append(cost)\n",
    "                fathers[father_i] = min(father_i,father_j)\n",
    "                fathers[father_j] = min(father_i,father_j)\n",
    "                ans+=cost\n",
    "            \n",
    "            if len(used_edges) == num_node:\n",
    "                break\n",
    "        \n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        self.parent = list(range(n+1))\n",
    "        self.cot = n+1\n",
    "        res = 0\n",
    "        # regard well as a node 0, being a water source\n",
    "        for i in range(n):\n",
    "            pipes.append([0,i+1,wells[i]])\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        idx_p = 0\n",
    "        while not self.cot==1:\n",
    "            cur = pipes[idx_p]\n",
    "            if not self.find(cur[0])==self.find(cur[1]):\n",
    "                self.union(cur[0], cur[1])\n",
    "                res += cur[2]\n",
    "            idx_p += 1\n",
    "        return res\n",
    "\n",
    "    def find(self, x):\n",
    "        while not x==self.parent[x]:\n",
    "            x = self.parent[x]\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        rootX = self.find(x)\n",
    "        rootY = self.find(y)\n",
    "        if rootX==rootY:\n",
    "            return\n",
    "        self.parent[rootX] = rootY\n",
    "        self.cot -= 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UF:\n",
    "    def __init__(self, n):\n",
    "        self.parent = list(range(n))\n",
    "        self.rank = [1] * n\n",
    "\n",
    "    def find(self, n):\n",
    "        while n != self.parent[n]:\n",
    "            self.parent[n] = self.parent[self.parent[n]]\n",
    "            n = self.parent[n]\n",
    "        return n\n",
    "    \n",
    "    def union(self, n1, n2):\n",
    "        p1, p2 = self.find(n1), self.find(n2)\n",
    "        if self.rank[p1] > self.rank[p2]:\n",
    "            self.parent[p2] = p1\n",
    "            self.rank[p1] += self.rank[p2]\n",
    "        else:\n",
    "            self.parent[p1] = p2\n",
    "            self.rank[p2] += self.rank[p1]\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        for v in range(1, n + 1):\n",
    "            pipes.append([0, v, wells[v - 1]])\n",
    "        \n",
    "        pipes.sort(key = lambda x: x[2])\n",
    "        uf, res, edge = UF(n + 1), 0, 0\n",
    "        for u, v, c in pipes:\n",
    "            if uf.find(u) != uf.find(v):\n",
    "                uf.union(u, v)\n",
    "                edge += 1\n",
    "                res += c\n",
    "            if edge == n:\n",
    "                break\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        self.parent = list(range(n+1))\n",
    "        self.cot = n+1\n",
    "        res = 0\n",
    "        # regard well as a node 0, being a water source\n",
    "        for i in range(n):\n",
    "            pipes.append([0,i+1,wells[i]])\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        idx_p = 0\n",
    "        while not self.cot==1:\n",
    "            cur = pipes[idx_p]\n",
    "            if self.union(cur[0], cur[1]):\n",
    "                res += cur[2]\n",
    "            idx_p += 1\n",
    "        return res\n",
    "\n",
    "    def find(self, x):\n",
    "        while not x==self.parent[x]:\n",
    "            self.parent[x] = self.parent[self.parent[x]]\n",
    "            x = self.parent[x]\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        rootX = self.find(x)\n",
    "        rootY = self.find(y)\n",
    "        if rootX==rootY:\n",
    "            return False\n",
    "        self.parent[rootX] = rootY\n",
    "        self.cot -= 1\n",
    "        return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def find(self, x):\n",
    "        if self.pa[x] != x:\n",
    "            self.pa[x] = self.find(self.pa[x])\n",
    "        return self.pa[x]\n",
    "\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        self.pa = list(range(n + 1))\n",
    "\n",
    "        for i, x in enumerate(wells):\n",
    "            pipes.append((0, i + 1, x))\n",
    "        pipes.sort(key = lambda x: x[2])\n",
    "        ans = 0\n",
    "        cnt = 0\n",
    "        for u, v, cost in pipes:\n",
    "            u, v = self.find(u), self.find(v)\n",
    "            if u != v:\n",
    "                self.pa[u] = v\n",
    "                ans += cost\n",
    "                cnt += 1\n",
    "            if cnt == n:\n",
    "                return ans\n",
    "        return -1\n",
    "\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UFS():                  #并查集类\n",
    "  def __init__(self,n):\n",
    "    self.n=n\n",
    "    self.parent=[0]*self.n                   #并查集存储结构\n",
    "    self.rnk=[-1]*self.n                    #存储结点的秩(近似于高度)\n",
    "  def Init(self):                               #并查集初始化\n",
    "    for i in range(0,self.n):\n",
    "      self.parent[i]=i\n",
    "      self.rnk[i]=0\n",
    "  def Find(self,x):         #递归算法：并查集中查找x结点的根结点\n",
    "    if x!=self.parent[x]:\n",
    "      self.parent[x]=self.Find(self.parent[x])        #路径压缩\n",
    "    return self.parent[x]\n",
    "  def Union(self,rx,ry):                        #并查集中x和y的两个集合的合并\n",
    "    if self.rnk[rx]<self.rnk[ry]:\n",
    "      self.parent[rx]=ry                 #rx结点作为ry的孩子\n",
    "    else:\n",
    "      if self.rnk[rx]==self.rnk[ry]:          #秩相同，合并后rx的秩增1\n",
    "        self.rnk[rx]+=1\n",
    "      self.parent[ry]=rx                  #ry结点作为rx的孩子\n",
    "\n",
    "class Solution:\n",
    "  def minCostToSupplyWater(self, n:int,wells:List[int],pipes:List[List[int]])->int:\n",
    "    for i in range(1,n+1):     #建立超级源点0\n",
    "      pipes.append([0,i,wells[i-1]])\n",
    "    return self.Kruskal(n+1,pipes)\n",
    "  def Kruskal(self,n,E):         #Kruskal算法\n",
    "    ufs=UFS(n) \n",
    "    E.sort(key=itemgetter(2))           #按边权值递增排序\n",
    "    ans=0\n",
    "    ufs.Init()                    #初始化并查集\n",
    "    k,j=0,0                           #k表示当前构造生成树的边数\n",
    "    while k<n-1:                        #生成的边数小于n-1时循环\n",
    "      u1,v1=E[j][0],E[j][1]               #取一条边(u1,v1)\n",
    "      sn1,sn2=ufs.Find(u1),ufs.Find(v1)   #两个顶点所属的集合编号\n",
    "      if sn1!=sn2:                        #添加该边不会构成回路\n",
    "        ans+=E[j][2]      #产生最小生成树的一条边\n",
    "        k+=1                #生成边数增1\n",
    "        ufs.Union(sn1,sn2)        #将sn1和sn2两个顶点合并\n",
    "      j+=1                #遍历下一条边\n",
    "    return ans\n",
    "\n",
    "       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UF:\n",
    "    def __init__(self,n):\n",
    "        self.fa = list(range(n))\n",
    "        self.sz = [1] * n\n",
    "        self.n_comp = n\n",
    "    def find(self,x):\n",
    "        if x!=self.fa[x]:\n",
    "            self.fa[x] = self.find(self.fa[x])\n",
    "        return self.fa[x]\n",
    "    def union(self,x,y):\n",
    "        x = self.find(x)\n",
    "        y = self.find(y)\n",
    "        if x==y:\n",
    "            return False\n",
    "        if self.sz[x] < self.sz[y]:\n",
    "            y,x = x,y\n",
    "        self.fa[y] = x\n",
    "        self.sz[x] + self.sz[y]\n",
    "        self.n_comp -= 1\n",
    "        return True\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        uf = UF(n+1)\n",
    "        for i,x in enumerate(wells):\n",
    "            pipes.append([0,i+1,x])\n",
    "\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        \n",
    "        ans = 0\n",
    "        for x,y,w in pipes:\n",
    "            if  uf.union(x,y):\n",
    "               ans += w \n",
    "\n",
    "            if uf.n_comp == 1:\n",
    "                break\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        self.parent = list(range(n+1))\n",
    "        self.cot = n+1\n",
    "        res = 0\n",
    "        # regard well as a node 0, being a water source\n",
    "        for i in range(n):\n",
    "            pipes.append([0,i+1,wells[i]])\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        idx_p = 0\n",
    "        while not self.cot==1:\n",
    "            cur = pipes[idx_p]\n",
    "            if self.union(cur[0], cur[1]):\n",
    "                res += cur[2]\n",
    "            idx_p += 1\n",
    "        return res\n",
    "\n",
    "    def find(self, x):\n",
    "        while not x==self.parent[x]:\n",
    "            x = self.parent[x]\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        rootX = self.find(x)\n",
    "        rootY = self.find(y)\n",
    "        if rootX==rootY:\n",
    "            return False\n",
    "        self.parent[rootX] = rootY\n",
    "        self.cot -= 1\n",
    "        return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        for i in range(n):\n",
    "            pipes.append([i+1,0,wells[i]])\n",
    "        ans = 0\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        father = [_ for _ in range(n+1)]\n",
    "        def find(x):\n",
    "            if father[x] == x:\n",
    "                return x\n",
    "            father[x] = find(father[x])\n",
    "            return father[x]\n",
    "        \n",
    "        def join(x,y):\n",
    "            father[find(x)] = find(y)\n",
    "        \n",
    "        index = 0\n",
    "        for i in range(n):\n",
    "            while find(pipes[index][0]) == find(pipes[index][1]):\n",
    "                index += 1\n",
    "            join(pipes[index][0], pipes[index][1])\n",
    "            ans += pipes[index][2]\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int) -> None:\n",
    "        self.root = [i for i in range(n)]\n",
    "        self.size = [1] * n\n",
    "        self.part = n\n",
    "        return\n",
    "\n",
    "    def find(self, x):\n",
    "        lst = []\n",
    "        while x != self.root[x]:\n",
    "            lst.append(x)\n",
    "            # 在查询的时候合并到顺带直接根节点\n",
    "            x = self.root[x]\n",
    "        for w in lst:\n",
    "            self.root[w] = x\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        root_x = self.find(x)\n",
    "        root_y = self.find(y)\n",
    "        if root_x == root_y:\n",
    "            return False\n",
    "        if self.size[root_x] >= self.size[root_y]:\n",
    "            root_x, root_y = root_y, root_x\n",
    "        self.root[root_x] = root_y\n",
    "        self.size[root_y] += self.size[root_x]\n",
    "        # 将非根节点的秩赋0\n",
    "        self.size[root_x] = 0\n",
    "        self.part -= 1\n",
    "        return True\n",
    "\n",
    "    def is_connected(self, x, y):\n",
    "        return self.find(x) == self.find(y)\n",
    "\n",
    "    def get_root_part(self):\n",
    "        # 获取每个根节点对应的组\n",
    "        part = defaultdict(list)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            part[self.find(i)].append(i)\n",
    "        return part\n",
    "\n",
    "    def get_root_size(self):\n",
    "        # 获取每个根节点对应的组大小\n",
    "        size = defaultdict(int)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            size[self.find(i)] = self.size[self.find(i)]\n",
    "        return size\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        uf =UnionFind(n+1)\n",
    "        for i in range(n):\n",
    "            pipes.append([0, i+1, wells[i]])\n",
    "        \n",
    "        pipes.sort(key=lambda it: it[2])\n",
    "        ans = 0\n",
    "        for i, j, w in pipes:\n",
    "            if uf.union(i, j):\n",
    "                ans += w\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        root = list(range(n + 1))\n",
    "        def Find(x):\n",
    "            if x != root[x]:\n",
    "                root[x] = Find(root[x])\n",
    "            return root[x]\n",
    "        def Union(x,y):\n",
    "            root[x] = y\n",
    "\n",
    "        for idx,well in enumerate(wells):\n",
    "            pipes.append([idx + 1,0,well])\n",
    "        pipes.sort(key = lambda x : x[-1])\n",
    "        print(pipes)\n",
    "        min_cost = 0\n",
    "        cnt = 0\n",
    "        for x,y,c in pipes:\n",
    "            rx,ry = Find(x),Find(y)\n",
    "            if rx != ry:\n",
    "                Union(rx,ry)\n",
    "                min_cost += c\n",
    "        return min_cost\n",
    "\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        p={i:i for i in range(n+1)}\n",
    "        for i,v in enumerate(wells):\n",
    "            pipes.append([0,i+1,v])\n",
    "\n",
    "        def find(x):\n",
    "            if p[x]==x:\n",
    "                return p[x]\n",
    "            p[x]=find(p[x])\n",
    "            return p[x]\n",
    "        def union(x,y):\n",
    "            p[find(x)]=find(y)\n",
    "        cnt=0\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        for s,e,c in pipes:\n",
    "            a=find(s)\n",
    "            b=find(e)\n",
    "            if a!=b:\n",
    "                union(s,e)\n",
    "                cnt+=c\n",
    "        return cnt \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        # 增加虚拟节点\n",
    "        # 新的边为对应节点建筑水井的代价\n",
    "        # 构建最小生成树\n",
    "        # 不知道几棵树-->想办法连接成为一棵树-->增加虚拟节点\n",
    "        Tree=[i for i in range(n+2)]\n",
    "        def find_root(idx):\n",
    "            if idx==Tree[idx]:\n",
    "                return idx\n",
    "            while idx!=Tree[idx]:\n",
    "                Tree[idx]=Tree[Tree[idx]]\n",
    "                idx=Tree[idx]\n",
    "            return idx\n",
    "        for i in range(len(wells)):\n",
    "            pipes.append([i+1, n+1, wells[i]])\n",
    "        pipes.sort(key=lambda x:x[2])\n",
    "        val=0\n",
    "        for pipe in pipes:\n",
    "            root_a=find_root(pipe[0])\n",
    "            root_b=find_root(pipe[1])\n",
    "            if root_a!=root_b:\n",
    "                Tree[root_a]=root_b\n",
    "                val+=pipe[2]\n",
    "        return val\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        edges = []\n",
    "        for index, w in enumerate(wells):\n",
    "            n_index = index+1\n",
    "            edges.append([0, n_index, w])\n",
    "        \n",
    "        pipes_cost = {}\n",
    "        for p in pipes:\n",
    "            s, t, v = p\n",
    "            if (s, t) not in pipes_cost:\n",
    "                pipes_cost[(s, t)] = v\n",
    "            else:\n",
    "                pipes_cost[(s, t)] = min(pipes_cost[(s, t)], v)\n",
    "        \n",
    "        for p, cost in pipes_cost.items():\n",
    "            edges.append([p[0], p[1], cost])\n",
    "        \n",
    "        edges.sort(key=lambda x:x[2])\n",
    "\n",
    "        parent = [i for i in range(n+1)]\n",
    "        def find(i):\n",
    "            if parent[i] != i:\n",
    "                parent[i] = find(parent[i])\n",
    "            return parent[i]\n",
    "        \n",
    "        def union(i, j):\n",
    "            irep, jrep = find(i), find(j)\n",
    "            if irep == jrep:\n",
    "                return\n",
    "            else:\n",
    "                parent[irep] = jrep\n",
    "        \n",
    "        total_cost = 0\n",
    "        for e in edges:\n",
    "            s, t, v = e\n",
    "            if find(s) != find(t):\n",
    "                union(s, t)\n",
    "                total_cost += v\n",
    "\n",
    "            if len(set(parent)) == 1:\n",
    "                break\n",
    "        \n",
    "        return total_cost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        # 邻接表表示的双向图\n",
    "       graph = defaultdict(list)\n",
    "\n",
    "       # 添加索引为0的虚拟顶点。\n",
    "       # 然后在每一栋房子上加一条边，按成本加权\n",
    "       for index, cost in enumerate(wells):\n",
    "           graph[0].append((cost, index + 1))\n",
    "\n",
    "       # 将双向边添加到图中\n",
    "       for house_1, house_2, cost in pipes:\n",
    "           graph[house_1].append((cost, house_2))\n",
    "           graph[house_2].append((cost, house_1))\n",
    "\n",
    "       # 用于维护已添加到的所有顶点的集合\n",
    "       # 最终的 MST(最小生成树)从顶点 0 开始。\n",
    "       mst_set = set([0])\n",
    "\n",
    "       # 堆维护要访问的边的顺序，\n",
    "       # 从起源于顶点0的边开始。\n",
    "       # 注意：我们能够从任意节点开始。\n",
    "       heapq.heapify(graph[0])\n",
    "       edges_heap = graph[0]\n",
    "\n",
    "       total_cost = 0\n",
    "       while len(mst_set) < n + 1:\n",
    "           cost, next_house = heapq.heappop(edges_heap)\n",
    "           if next_house not in mst_set:\n",
    "               # 将新顶点添加到集合中\n",
    "               mst_set.add(next_house)\n",
    "               total_cost += cost\n",
    "               # 在下一轮扩大候选边的选择范围\n",
    "               for new_cost, neighbor_house in graph[next_house]:\n",
    "                   if neighbor_house not in mst_set:\n",
    "                       heapq.heappush(edges_heap, (new_cost, neighbor_house))\n",
    "\n",
    "       return total_cost\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import heapq\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        graph = defaultdict(list)\n",
    "        for index, cost in enumerate(wells):\n",
    "            graph[0].append((cost, index+1))\n",
    "            # graph[index+1].append((cost, 0))\n",
    "        for house_1,house_2,cost in pipes:\n",
    "            graph[house_1].append((cost, house_2))\n",
    "            graph[house_2].append((cost, house_1))\n",
    "        mst_set = set([0])\n",
    "        heapq.heapify(graph[0])\n",
    "        edges_heap = graph[0]\n",
    "\n",
    "        total_cost = 0\n",
    "        while len(mst_set)<n+1:\n",
    "            cost, next_house = heapq.heappop(edges_heap)\n",
    "            if next_house not in mst_set:\n",
    "                mst_set.add(next_house)\n",
    "                total_cost += cost\n",
    "                for new_cost,neighbor_house in graph[next_house]:\n",
    "                    if neighbor_house not in mst_set:\n",
    "                        heapq.heappush(edges_heap, (new_cost,neighbor_house))\n",
    "                        \n",
    "        return total_cost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import heapq\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        graph = defaultdict(list)\n",
    "        for index, cost in enumerate(wells):\n",
    "            graph[0].append((cost, index+1))\n",
    "            # graph[index+1].append((cost, 0))\n",
    "        for house_1,house_2,cost in pipes:\n",
    "            graph[house_1].append((cost, house_2))\n",
    "            graph[house_2].append((cost, house_1))\n",
    "        mst_set = set([0])\n",
    "        heapq.heapify(graph[0])\n",
    "        edges_heap = graph[0]\n",
    "\n",
    "        res = 0\n",
    "        while len(mst_set)<n+1:\n",
    "            cost, nextN = heapq.heappop(edges_heap)\n",
    "            if nextN not in mst_set:\n",
    "                res += cost\n",
    "                mst_set.add(nextN)\n",
    "                for c,neighbor in graph[nextN]:\n",
    "                    if neighbor not in mst_set:\n",
    "                        heapq.heappush(edges_heap, (c,neighbor))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "   def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "\n",
    "       # 邻接表表示的双向图\n",
    "       graph = defaultdict(list)\n",
    "\n",
    "       # 添加索引为0的虚拟顶点。\n",
    "       # 然后在每一栋房子上加一条边，按成本加权\n",
    "       for index, cost in enumerate(wells):\n",
    "           graph[0].append((cost, index + 1))\n",
    "\n",
    "       # 将双向边添加到图中\n",
    "       for house_1, house_2, cost in pipes:\n",
    "           graph[house_1].append((cost, house_2))\n",
    "           graph[house_2].append((cost, house_1))\n",
    "\n",
    "       # 用于维护已添加到的所有顶点的集合\n",
    "       # 最终的 MST(最小生成树)从顶点 0 开始。\n",
    "       mst_set = set([0])\n",
    "\n",
    "       # 堆维护要访问的边的顺序，\n",
    "       # 从起源于顶点0的边开始。\n",
    "       # 注意：我们能够从任意节点开始。\n",
    "       heapq.heapify(graph[0])\n",
    "       edges_heap = graph[0]\n",
    "\n",
    "       total_cost = 0\n",
    "       while len(mst_set) < n + 1:\n",
    "           cost, next_house = heapq.heappop(edges_heap)\n",
    "           if next_house not in mst_set:\n",
    "               # 将新顶点添加到集合中\n",
    "               mst_set.add(next_house)\n",
    "               total_cost += cost\n",
    "               # 在下一轮扩大候选边的选择范围\n",
    "               for new_cost, neighbor_house in graph[next_house]:\n",
    "                   if neighbor_house not in mst_set:\n",
    "                       heapq.heappush(edges_heap, (new_cost, neighbor_house))\n",
    "\n",
    "       return total_cost"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "   def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "\n",
    "       # 邻接表表示的双向图\n",
    "       graph = defaultdict(list)\n",
    "\n",
    "       # 添加索引为0的虚拟顶点。\n",
    "       # 然后在每一栋房子上加一条边，按成本加权\n",
    "       for index, cost in enumerate(wells):\n",
    "           graph[0].append((cost, index + 1))\n",
    "\n",
    "       # 将双向边添加到图中\n",
    "       for house_1, house_2, cost in pipes:\n",
    "           graph[house_1].append((cost, house_2))\n",
    "           graph[house_2].append((cost, house_1))\n",
    "\n",
    "       # 用于维护已添加到的所有顶点的集合\n",
    "       # 最终的 MST(最小生成树)从顶点 0 开始。\n",
    "       mst_set = set([0])\n",
    "\n",
    "       # 堆维护要访问的边的顺序，\n",
    "       # 从起源于顶点0的边开始。\n",
    "       # 注意：我们能够从任意节点开始。\n",
    "       heapq.heapify(graph[0])\n",
    "       edges_heap = graph[0]\n",
    "\n",
    "       total_cost = 0\n",
    "       while len(mst_set) < n + 1:\n",
    "           cost, next_house = heapq.heappop(edges_heap)\n",
    "           if next_house not in mst_set:\n",
    "               # 将新顶点添加到集合中\n",
    "               mst_set.add(next_house)\n",
    "               total_cost += cost\n",
    "               # 在下一轮扩大候选边的选择范围\n",
    "               for new_cost, neighbor_house in graph[next_house]:\n",
    "                   if neighbor_house not in mst_set:\n",
    "                       heapq.heappush(edges_heap, (new_cost, neighbor_house))\n",
    "\n",
    "       return total_cost\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import heapq\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        graph = defaultdict(list)\n",
    "        for index, cost in enumerate(wells):\n",
    "            graph[0].append((cost, index+1))\n",
    "            # graph[index+1].append((cost, 0))\n",
    "        for house_1,house_2,cost in pipes:\n",
    "            graph[house_1].append((cost, house_2))\n",
    "            graph[house_2].append((cost, house_1))\n",
    "        mst_set = set([0])\n",
    "\n",
    "        edges_heap = graph[0]\n",
    "        heapq.heapify(edges_heap)\n",
    "\n",
    "        res = 0\n",
    "        while len(mst_set)<n+1:\n",
    "            cost, nextN = heapq.heappop(edges_heap)\n",
    "            if nextN not in mst_set:\n",
    "                res += cost\n",
    "                mst_set.add(nextN)\n",
    "                for c,neighbor in graph[nextN]:\n",
    "                    if neighbor not in mst_set:\n",
    "                        heapq.heappush(edges_heap, (c,neighbor))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import heapq\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        graph = defaultdict(list)\n",
    "        for index, cost in enumerate(wells):\n",
    "            graph[0].append((cost, index+1))\n",
    "            # graph[index+1].append((cost, 0))\n",
    "        for house_1,house_2,cost in pipes:\n",
    "            graph[house_1].append((cost, house_2))\n",
    "            graph[house_2].append((cost, house_1))\n",
    "        mst_set = set([0])\n",
    "\n",
    "        edges_heap = graph[0]\n",
    "        heapq.heapify(edges_heap)\n",
    "\n",
    "        res = 0\n",
    "        while len(mst_set)<n+1:\n",
    "            cost, nextN = heapq.heappop(edges_heap)\n",
    "            if nextN not in mst_set:\n",
    "                res += cost\n",
    "                mst_set.add(nextN)\n",
    "                for c,neighbor in graph[nextN]:\n",
    "                    if neighbor not in mst_set:\n",
    "                        heapq.heappush(edges_heap, (c,neighbor))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "   def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "\n",
    "       # 邻接表表示的双向图\n",
    "       graph = defaultdict(list)\n",
    "\n",
    "       # 添加索引为0的虚拟顶点。\n",
    "       # 然后在每一栋房子上加一条边，按成本加权\n",
    "       for index, cost in enumerate(wells):\n",
    "           graph[0].append((cost, index + 1))\n",
    "\n",
    "       # 将双向边添加到图中\n",
    "       for house_1, house_2, cost in pipes:\n",
    "           graph[house_1].append((cost, house_2))\n",
    "           graph[house_2].append((cost, house_1))\n",
    "\n",
    "       # 用于维护已添加到的所有顶点的集合\n",
    "       # 最终的 MST(最小生成树)从顶点 0 开始。\n",
    "       mst_set = set([0])\n",
    "\n",
    "       # 堆维护要访问的边的顺序，\n",
    "       # 从起源于顶点0的边开始。\n",
    "       # 注意：我们能够从任意节点开始。\n",
    "       heapq.heapify(graph[0])\n",
    "       edges_heap = graph[0]\n",
    "\n",
    "       total_cost = 0\n",
    "       while len(mst_set) < n + 1:\n",
    "           cost, next_house = heapq.heappop(edges_heap)\n",
    "           if next_house not in mst_set:\n",
    "               # 将新顶点添加到集合中\n",
    "               mst_set.add(next_house)\n",
    "               total_cost += cost\n",
    "               # 在下一轮扩大候选边的选择范围\n",
    "               for new_cost, neighbor_house in graph[next_house]:\n",
    "                   if neighbor_house not in mst_set:\n",
    "                       heapq.heappush(edges_heap, (new_cost, neighbor_house))\n",
    "\n",
    "       return total_cost\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        ref = defaultdict(list)\n",
    "        for i, j, k in pipes:\n",
    "            ref[i].append((k, j))\n",
    "            ref[j].append((k, i))\n",
    "        heaq = []\n",
    "        for i in range(n):\n",
    "            heappush(heaq, (wells[i], i + 1))\n",
    "        visited = set()\n",
    "        res = 0\n",
    "        while heaq:\n",
    "            cost, node = heappop(heaq)\n",
    "            if node in visited:\n",
    "                continue\n",
    "            visited.add(node)\n",
    "            res += cost\n",
    "            for c, n in ref[node]:\n",
    "                if n not in visited:\n",
    "                    heappush(heaq, (c, n))\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        #解题思路：\n",
    "        #一个房子，要么连接水管，要么铺管道，不知道如何处理\n",
    "        #技巧：在图中添加一个虚拟顶点0，0与各边相连，其权值为水井的成本\n",
    "\n",
    "        edges = [[] for i in range(n + 1)]\n",
    "        for a,b,cost in pipes:\n",
    "            edges[a].append((cost,b))\n",
    "            edges[b].append((cost,a))\n",
    "        for i,cost in enumerate(wells):\n",
    "            edges[0].append((cost,i + 1))\n",
    "\n",
    "        visited = set()\n",
    "        out_edges = [(0,0)] #从顶点0开始\n",
    "        res = 0\n",
    "        while out_edges and len(visited) != (n + 1):\n",
    "            cost,city = heapq.heappop(out_edges)\n",
    "            if city not in visited:\n",
    "                visited.add(city)\n",
    "                res += cost\n",
    "                for next_cost,next_city in edges[city]:\n",
    "                    heapq.heappush(out_edges,(next_cost,next_city))\n",
    "        return res if len(visited) == (n + 1) else -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        #解题思路：\n",
    "        #一个房子，要么连接水管，要么铺管道，不知道如何处理\n",
    "        #技巧：在图中添加一个虚拟顶点0，0与各边相连，其权值为水井的成本\n",
    "\n",
    "        edges = [[] for i in range(n + 1)]\n",
    "        for a,b,cost in pipes:\n",
    "            edges[a].append((cost,b))\n",
    "            edges[b].append((cost,a))\n",
    "        for i,cost in enumerate(wells):\n",
    "            edges[0].append((cost,i + 1))\n",
    "\n",
    "        visited = set()\n",
    "        out_edges = [(0,0)] #从顶点0开始\n",
    "        res = 0\n",
    "        while out_edges and len(visited) != (n + 1):\n",
    "            cost,city = heapq.heappop(out_edges)\n",
    "            if city not in visited:\n",
    "                visited.add(city)\n",
    "                res += cost\n",
    "                for next_cost,next_city in edges[city]:\n",
    "                    heapq.heappush(out_edges,(next_cost,next_city))\n",
    "        return res if len(visited) == (n + 1) else -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        edges = [[] for _ in range(n + 1)]\n",
    "        for a, b, cost in pipes:\n",
    "            edges[a].append((b, cost))\n",
    "            edges[b].append((a, cost))\n",
    "        for i, v in enumerate(wells):\n",
    "            edges[0].append((i + 1, v))\n",
    "        \n",
    "        \n",
    "        tree = []\n",
    "        vis = [False] * (n + 1)\n",
    "        h = [(0, 0)]\n",
    "        arr = [i for i in range(n + 1)]\n",
    "        res = 0\n",
    "        while h:\n",
    "            cost, p = heapq.heappop(h)\n",
    "            \n",
    "            if not vis[p]:\n",
    "                vis[p] = True\n",
    "                tree.append(p)\n",
    "                res += cost\n",
    "                \n",
    "                for next_pos, next_cost in edges[p]:\n",
    "                    if not vis[next_pos]:\n",
    "                        heapq.heappush(h, (next_cost, next_pos))\n",
    "        return res\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        edges = [[] for _ in range(n + 1)]\n",
    "        for a, b, cost in pipes:\n",
    "            edges[a].append((b, cost))\n",
    "            edges[b].append((a, cost))\n",
    "        for i, v in enumerate(wells):\n",
    "            edges[0].append((i + 1, v))\n",
    "        \n",
    "        \n",
    "        tree = []\n",
    "        vis = [False] * (n + 1)\n",
    "        h = [(0, 0)]\n",
    "        arr = [i for i in range(n + 1)]\n",
    "        res = 0\n",
    "        while h:\n",
    "            cost, p = heapq.heappop(h)\n",
    "            \n",
    "            if not vis[p]:\n",
    "                vis[p] = True\n",
    "                tree.append(p)\n",
    "                res += cost\n",
    "                \n",
    "                for next_pos, next_cost in edges[p]:\n",
    "                    if not vis[next_pos]:\n",
    "                        heapq.heappush(h, (next_cost, next_pos))\n",
    "        return res\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "from sortedcontainers import SortedList\n",
    "from collections import deque, defaultdict\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        \"\"\"\n",
    "\n",
    "        :param n:\n",
    "        :param wells:\n",
    "        :param pipes:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        anslist = SortedList(key=lambda x: x[0])\n",
    "        idxmap = {}\n",
    "        for i in range(n):\n",
    "            anslist.add((wells[i], i))\n",
    "            idxmap[i] = wells[i]\n",
    "        ret = []\n",
    "        m1 = defaultdict(list)\n",
    "        for i in pipes:\n",
    "            m1[i[0]-1].append((i[1]-1, i[2]))\n",
    "            m1[i[1]-1].append((i[0]-1, i[2]))\n",
    "        processed = set()\n",
    "        while anslist:\n",
    "            val, idx = anslist.pop(0)\n",
    "            ret.append((val, idx))\n",
    "            processed.add(idx)\n",
    "            for i in m1.get(idx, []):\n",
    "                newxx = i[1]\n",
    "                if i[0] in processed:\n",
    "                    continue\n",
    "                if newxx < idxmap[i[0]]:\n",
    "\n",
    "                    anslist.remove((idxmap[i[0]], i[0]))\n",
    "                    idxmap[i[0]] = newxx\n",
    "                    anslist.add((newxx, i[0]))\n",
    "        return sum([i[0] for i in ret])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "from sortedcontainers import SortedList\n",
    "from collections import deque, defaultdict\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:\n",
    "        \"\"\"\n",
    "\n",
    "        :param n:\n",
    "        :param wells:\n",
    "        :param pipes:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        anslist = SortedList(key=lambda x: x[0])\n",
    "        idxmap = {}\n",
    "        for i in range(n):\n",
    "            anslist.add((wells[i], i))\n",
    "            idxmap[i] = wells[i]\n",
    "        ret = []\n",
    "        m1 = defaultdict(list)\n",
    "        for i in pipes:\n",
    "            m1[i[0]-1].append((i[1]-1, i[2]))\n",
    "            m1[i[1]-1].append((i[0]-1, i[2]))\n",
    "        processed = set()\n",
    "        while anslist:\n",
    "            val, idx = anslist.pop(0)\n",
    "            ret.append((val, idx))\n",
    "            processed.add(idx)\n",
    "            for i in m1.get(idx, []):\n",
    "                newxx = i[1]\n",
    "                if i[0] in processed:\n",
    "                    continue\n",
    "                if newxx < idxmap[i[0]]:\n",
    "                    anslist.remove((idxmap[i[0]], i[0]))\n",
    "                    idxmap[i[0]] = newxx\n",
    "                    anslist.add((newxx, i[0]))\n",
    "        return sum([i[0] for i in ret])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minCostToSupplyWater(self, n:int,wells:List[int],pipes:List[List[int]])->int:\n",
    "        if n==1:return wells[n-1]\n",
    "        adj=[[] for _ in range(n+1)] \n",
    "        for e in pipes:\n",
    "          a,b,w=e[0],e[1],e[2]\n",
    "          adj[a].append([b,w])\n",
    "          adj[b].append([a,w])\n",
    "        for i in range(1,n+1):      #建立超级源点0\n",
    "            adj[0].append([i,wells[i-1]])\n",
    "            adj[i].append([0,wells[i-1]])\n",
    "        U=[False for _ in range(n+1)]\n",
    "        ans=0          #存放答案\n",
    "        k=0        #累计选择的边数\n",
    "        pqu=[]          #小根堆，结点类型为[w,vno]\n",
    "        U[0]=True\n",
    "        for e in adj[0]:heapq.heappush(pqu,[e[1],e[0]])\n",
    "        while pqu:\n",
    "            [w,v]=heapq.heappop(pqu)       \t\t\t#出队\n",
    "            if U[v]==False:\n",
    "                U[v]=True\n",
    "                ans+=w\n",
    "                k+=1\n",
    "                if k==n:return ans\n",
    "                for e in adj[v]:heapq.heappush(pqu,[e[1],e[0]])\n",
    "        return -1\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
